import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.weight_norm import weight_norm
class Classifier(nn.Module):
    def __init__(self, dims):
        super(Classifier, self).__init__()

        layers = []
        for i in range(len(dims) - 2):
            in_dim = dims[i]
            out_dim = dims[i + 1]
            layers.append(weight_norm(nn.Linear(in_dim, out_dim), dim=None))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout())
        layers.append(weight_norm(nn.Linear(dims[-2], dims[-1]), dim=None))


        self.main = nn.Sequential(*layers)

    def forward(self, x):
        x = self.main(x)
        return F.log_softmax(x)